#include "cartographer/mapping/pose_extrapolator.h"
#include <algorithm>
#include "absl/memory/memory.h"
#include "glog/logging.h"

namespace cartographer
{
    namespace mapping
    {
        std::unique_ptr<PoseExtrapolator> PoseExtrapolator::InitializeWithImu(
            const common::Duration pose_queue_duration,
            const double imu_gravity_time_constant,
            const sensor::ImuData &imu_data)
        {
            auto extrapolator = absl::make_unique<PoseExtrapolator>(pose_queue_duration, imu_gravity_time_constant);
            extrapolator->AddImuData(imu_data);
            extrapolator->imu_tracker_ = absl::make_unique<ImuTracker>(imu_gravity_time_constant, imu_data.time);
            extrapolator->imu_tracker_->AddImuLinearAccelerationObservation(imu_data.linear_acceleration);
            extrapolator->imu_tracker_->AddImuAngularVelocityObservation(imu_data.angular_velocity);
            extrapolator->imu_tracker_->Advance(imu_data.time);
            extrapolator->AddPose(imu_data.time, transform::Rigid3d::Rotation(extrapolator->imu_tracker_->orientation()));
            return extrapolator;
        }

        // Returns the time of the last added pose or Time::min() if no pose was added yet.
        common::Time PoseExtrapolator::GetLastPoseTime() const
        {
            if (timed_pose_queue_.empty())
            {
                return common::Time::min();
            }
            return timed_pose_queue_.back().time;
        }

        Eigen::Vector3d PoseExtrapolator::GetLinearVelocity()
        {
            if (odometry_data_.size() < 2)
            {
                return linear_velocity_from_poses_;
            }
            return linear_velocity_from_odometry_;
        }

        common::Time PoseExtrapolator::GetLastExtrapolatedTime() const
        {
            if (!extrapolation_imu_tracker_)
            {
                return common::Time::min();
            }
            return extrapolation_imu_tracker_->time();
        }

        void PoseExtrapolator::PrintPose(const char *prefix, const TimedPose &pose)
        {
            const auto &trans = pose.pose.translation();
            const auto &rotation = transform::QuatToEuler(pose.pose.rotation());
            LOG(INFO) << prefix << " \n trans:" << VecToStr(trans)
                      << " rotat:" << VecToStr(rotation)
                      << " time:" << common::ToUniversal(pose.time) << "\n"
                      << " lin_vel: " << VecToStr(linear_velocity_from_poses_) << "\n\n";
        }

        PoseExtrapolator::PoseExtrapolator(const common::Duration pose_queue_duration, double imu_gravity_time_constant)
            : gravity_time_constant_(imu_gravity_time_constant),
              pose_queue_duration_(pose_queue_duration), // 两次预测位姿需要的最小时间间隔
              cached_extrapolated_pose_{common::Time::min(), transform::Rigid3d::Identity()}
        {
            // poseIntegrator_ = std::make_shared<PoseIntegrator>();
        }

        void PoseExtrapolator::UpdateVelocitiesFromPoses()
        {
            if (timed_pose_queue_.size() < 2)
            {
                // We need two poses to estimate velocities.
                return;
            }
            CHECK(!timed_pose_queue_.empty());
            const TimedPose &newest_timed_pose = timed_pose_queue_.back();
            const auto newest_time = newest_timed_pose.time;
            const TimedPose &oldest_timed_pose = timed_pose_queue_.front();
            const auto oldest_time = oldest_timed_pose.time;
            const double queue_delta = common::ToSeconds(newest_time - oldest_time);
            if (queue_delta < common::ToSeconds(pose_queue_duration_))
            {
                LOG(WARNING) << "Queue too short for velocity estimation. Queue duration: " << queue_delta << " s";
                return;
            }
            const transform::Rigid3d &newest_pose = newest_timed_pose.pose;
            const transform::Rigid3d &oldest_pose = oldest_timed_pose.pose;
            // calculate transform delta
            const auto &translation_delta = newest_pose.translation() - oldest_pose.translation();
            const auto &quaternion_delta = oldest_pose.rotation().inverse() * newest_pose.rotation();
            const auto &angle_delta = transform::RotationQuaternionToAngleAxisVector(quaternion_delta);
            // divide time delta into transform delta, then get the velocity
            angular_velocity_from_poses_ = angle_delta / queue_delta;
            linear_velocity_from_poses_ = translation_delta / queue_delta;
        }

        /*
         * 每当前端CSM计算姿态后，会调用该函数，此时点云匹配的姿态(timedpose)会被用于统一更新/对齐该类的三个ImuTracker
         */
        void PoseExtrapolator::AddPose(const common::Time time, const transform::Rigid3d &pose)
        {
            LOG(INFO) << "AddPose 1\n";
            if (imu_tracker_ == nullptr)
            {
                common::Time tracker_start = time;
                if (!imu_data_.empty())
                {
                    tracker_start = std::min(tracker_start, imu_data_.front().time);
                }
                imu_tracker_ = absl::make_unique<ImuTracker>(gravity_time_constant_, tracker_start);
            }
            auto timePose = TimedPose{time, pose};
            timed_pose_queue_.push_back(timePose);
            while (timed_pose_queue_.size() > 2 && timed_pose_queue_[1].time <= time - pose_queue_duration_)
            {
                timed_pose_queue_.pop_front();
            }
            LOG(INFO) << "AddPose 2\n";
            UpdateVelocitiesFromPoses();
            // PrintPose("AddPose", timePose);
            AdvanceImuTracker(time, imu_tracker_.get());
            TrimImuData();
            TrimOdometryData();
            LOG(INFO) << "AddPose end\n";
            odometry_imu_tracker_ = absl::make_unique<ImuTracker>(*imu_tracker_);
            extrapolation_imu_tracker_ = absl::make_unique<ImuTracker>(*imu_tracker_);
        }

        void PoseExtrapolator::AddImuData(const sensor::ImuData &data)
        {
            CHECK(timed_pose_queue_.empty() || data.time >= GetLastPoseTime());
            imu_data_.push_back(data);
            TrimImuData();
            LOG(INFO) << "AddImuData\n";
        }

        /*
         * odom的坐标系与激光匹配的坐标系不一致，需要将通过odom data估算出来的速度进行坐标变换。
         */
        void PoseExtrapolator::AddOdometryData(const sensor::OdometryData &odometry_data)
        {
            CHECK(timed_pose_queue_.empty() || odometry_data.time >= GetLastPoseTime());
            odometry_data_.push_back(odometry_data);
            TrimOdometryData();
            if (odometry_data_.size() < 2)
                return;
            // TODO(whess): Improve by using more than just the last two odometry poses.
            // Compute extrapolation in the tracking frame.
            const sensor::OdometryData &odometry_data_oldest = odometry_data_.front();
            const sensor::OdometryData &odometry_data_newest = odometry_data_.back();
            const double odometry_time_delta = common::ToSeconds(odometry_data_oldest.time - odometry_data_newest.time);
            const transform::Rigid3d odometry_pose_delta = odometry_data_newest.pose.inverse() * odometry_data_oldest.pose;
            angular_velocity_from_odometry_ = transform::RotationQuaternionToAngleAxisVector(odometry_pose_delta.rotation()) / odometry_time_delta;
            if (timed_pose_queue_.empty())
            {
                return;
            }
            // This moment, the newest odom must be later than the newest pose. So we can utilize the newest pose and
            // the rotation delta of odom to estimate the current pose in odometry frame.
            // linear velocity in tracking frame at newest odometry time
            const Eigen::Vector3d linear_velocity_in_tracking_frame = odometry_pose_delta.translation() / odometry_time_delta;
            const auto &quaternion_delta_odom = ExtrapolateRotation(odometry_data_newest.time, odometry_imu_tracker_.get());
            // orientation at newest odometry time
            const Eigen::Quaterniond orientation_odom = timed_pose_queue_.back().pose.rotation() * quaternion_delta_odom;
            linear_velocity_from_odometry_ = orientation_odom * linear_velocity_in_tracking_frame;
        }

        transform::Rigid3d PoseExtrapolator::ExtrapolateRobotPose(const common::Time time)
        {
            int poseSize = timed_pose_queue_.size();
            const auto &newest_pose = timed_pose_queue_.back();
            CHECK_GE(time, newest_pose.time);
            LOG(INFO) << "ExtrapolateRobotPose\n";
            if (cached_extrapolated_pose_.time != time)
            {
                const Eigen::Vector3d delta_trans = ExtrapolateTranslation(time);
                const Eigen::Vector3d translation = delta_trans + newest_pose.pose.translation();
                const Eigen::Quaterniond rotate_delta = ExtrapolateRotation(time, extrapolation_imu_tracker_.get());
                const Eigen::Quaterniond rotation = newest_pose.pose.rotation() * rotate_delta;
                cached_extrapolated_pose_ = TimedPose{time, transform::Rigid3d{translation, rotation}};
            }
            if (use_imu_integrate && poseSize > 2)
            {
                float duration = common::ToSeconds(time - newest_pose.time);
                LOG(INFO) << "[pose] ExtraRobotPose poseSize " << poseSize << " dt " << duration
                          << " " << common::ToUniversal(newest_pose.time)
                          << " -> " << common::ToUniversal(time) << "\n";
                if (time > newest_pose.time)
                {
                    int collectRes = 0;
                    int integrateRes = 0;
                    std::vector<sensor::ImuData> imuSet;
                    collectRes = CollectImuByTimeSpan(imu_data_, imuSet, newest_pose.time, time);
                    int dataSize = imuSet.size();
                    LOG(INFO) << "[pose] collectImuRes: " << collectRes << " size " << dataSize << "\n";
                    if (collectRes > 0 && dataSize >= 2)
                    {
                        // poseIntegrator_->linear_velocity_ = GetLinearVelocity();
                        // integrateRes = poseIntegrator_->Integrate(imuSet, newest_pose.time, time);
                        if (integrateRes > 0)
                        {
                            TimedPose imuPose;
                            // poseIntegrator_->Predict(newest_pose, time, imuPose);
                            PrintPose("ImuPose", imuPose);
                            PrintPose("ExtPose", cached_extrapolated_pose_);
                        }
                    }
                    LOG(INFO) << "[pose] integrateRes " << integrateRes << " collectRes " << collectRes << "\n";
                }
            }
            return cached_extrapolated_pose_.pose;
        }

        transform::Rigid3d PoseExtrapolator::ExtrapolatePose(const common::Time time)
        {
            const TimedPose &newest_timed_pose = timed_pose_queue_.back();
            CHECK_GE(time, newest_timed_pose.time);
            if (cached_extrapolated_pose_.time != time)
            {
                const Eigen::Vector3d delta_trans = ExtrapolateTranslation(time);
                const Eigen::Vector3d translation = delta_trans + newest_timed_pose.pose.translation();
                const Eigen::Quaterniond rotate_delta = ExtrapolateRotation(time, extrapolation_imu_tracker_.get());
                const Eigen::Quaterniond rotation = newest_timed_pose.pose.rotation() * rotate_delta;
                cached_extrapolated_pose_ = TimedPose{time, transform::Rigid3d{translation, rotation}};
            }
            return cached_extrapolated_pose_.pose;
        }

        Eigen::Quaterniond PoseExtrapolator::EstimateGravityOrientation(const common::Time time)
        {
            ImuTracker imu_tracker = *imu_tracker_;
            AdvanceImuTracker(time, &imu_tracker);
            return imu_tracker.orientation();
        }

        // imu data must be later than the last pose data.
        // 将IMU队列中的第二条数据的时间戳与timed_pose_queue_的最新时间戳进行比较，如果不比最新的timed_pose的时间新，则删除第一条数据。
        // 意思是IMU的数据频率应该比雷达数据频率高，在最近的一次SCAN匹配成功结束后还应该收到更新的IMU数据。TrimOdometryData的逻辑与此同理。
        void PoseExtrapolator::TrimImuData()
        {
            while (imu_data_.size() > 1 && !timed_pose_queue_.empty() && imu_data_[1].time <= GetLastPoseTime())
            {
                imu_data_.pop_front();
            }
        }

        void PoseExtrapolator::TrimOdometryData()
        {
            while (odometry_data_.size() > 2 && !timed_pose_queue_.empty() && odometry_data_[1].time <= GetLastPoseTime())
            {
                odometry_data_.pop_front();
            }
        }

        /*
         * 该函数的目的就是更新ImuTracker中的orientation，从而估计最新的旋转量。
         * 使用IMU队列在time时间后面的IMU数据(线加速度和角速度)更新ImuTracker中的orientation, gravity_vector_，time_
         * 和imu_angular_velocity_。IMU数据中的线加速度不会被用来估计线性平移，它只会被用来更新gravity_vector_。
         */
        void PoseExtrapolator::AdvanceImuTracker(const common::Time time, ImuTracker *const imu_tracker) const
        {
            CHECK_GE(time, imu_tracker->time());
            if (imu_data_.empty() || time < imu_data_.front().time)
            {
                // There is no IMU data until 'time', so we advance the ImuTracker and use
                // the angular velocities from poses and fake gravity to help 2D stability.
                imu_tracker->Advance(time);
                // 当不使用IMU数据时，angular_velocity_from_poses_或angular_velocity_from_odometry_数据被传递给imu_tracker。
                const auto &angular_velocity = odometry_data_.size() < 2 ? angular_velocity_from_poses_ : angular_velocity_from_odometry_;
                imu_tracker->AddImuLinearAccelerationObservation(Eigen::Vector3d::UnitZ());
                imu_tracker->AddImuAngularVelocityObservation(angular_velocity);
                return;
            }
            if (imu_tracker->time() < imu_data_.front().time)
            {
                // Advance to the beginning of 'imu_data_'.
                imu_tracker->Advance(imu_data_.front().time);
            }
            auto it = std::lower_bound(
                imu_data_.begin(), imu_data_.end(), imu_tracker->time(),
                [](const sensor::ImuData &imu_data, const common::Time &time) {
                    return imu_data.time < time;
                });
            while (it != imu_data_.end() && it->time < time)
            {
                imu_tracker->Advance(it->time);
                imu_tracker->AddImuLinearAccelerationObservation(it->linear_acceleration);
                imu_tracker->AddImuAngularVelocityObservation(it->angular_velocity);
                ++it;
            }
            imu_tracker->Advance(time);
        }

        Eigen::Quaterniond PoseExtrapolator::ExtrapolateRotation(const common::Time time, ImuTracker *const imu_tracker) const
        {
            CHECK_GE(time, imu_tracker->time());
            // 先更新ImuTracker中的orientation，然后获取ImuTracker中的orientation
            AdvanceImuTracker(time, imu_tracker);
            const Eigen::Quaterniond last_orientation = imu_tracker_->orientation();
            return last_orientation.inverse() * imu_tracker->orientation();
        }

        /*
         * 如使用了odom数据，则使用linear_velocity_from_odometry_来估计平移，否则使用linear_velocity_from_poses_来估计平移。
         * 所以该类的平移估计的前提是假设机器人在短时间内是均速运动的
         */
        Eigen::Vector3d PoseExtrapolator::ExtrapolateTranslation(common::Time time)
        {
            const TimedPose &newest_timed_pose = timed_pose_queue_.back();
            const double delta_time = common::ToSeconds(time - newest_timed_pose.time);
            if (std::fabs(delta_time) > 10.0)
            {
                LOG(INFO) << "time: " << common::ToUniversal(time) << " " << common::ToUniversal(newest_timed_pose.time) << "\n";
                LOG(INFO) << "extra trans dt: " << delta_time << " size " << timed_pose_queue_.size() << "\n";
                Print("linear_velocity: ", linear_velocity_from_poses_);
            }
            return delta_time * GetLinearVelocity();
        }

        PoseExtrapolator::ExtrapolationResult PoseExtrapolator::ExtrapolatePosesWithGravity(const std::vector<common::Time> &times)
        {
            std::vector<transform::Rigid3f> poses;
            for (auto it = times.begin(); it != std::prev(times.end()); ++it)
            {
                poses.push_back(ExtrapolatePose(*it).cast<float>());
            }
            // 第二个参数是current_pose,将作为ScanMatch的初始位姿
            // return ExtrapolationResult{ poses, ExtrapolatePose(times.back()),GetLinearVelocity(),EstimateGravityOrientation(times.back()) };
            return ExtrapolationResult{poses, ExtrapolateRobotPose(times.back()), GetLinearVelocity(), EstimateGravityOrientation(times.back())};
        }

    } // namespace mapping
} // namespace cartographer
